import torch
import torch.nn as nn
import torch.optim as optim

# 定义模型
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(10, 64)  # 输入层到隐藏层
        self.fc2 = nn.Linear(64, 1)   # 隐藏层到输出层

    def forward(self, x):
        x = torch.relu(self.fc1(x))   # 激活函数
        return self.fc2(x)

# 初始化模型、损失函数和优化器
model = SimpleNN()
criterion = nn.MSELoss()  # 均方误差损失
optimizer = optim.Adam(model.parameters(), lr=0.001)  # Adam优化器

# 训练模型
for epoch in range(100):
    model.train()
    optimizer.zero_grad()
    outputs = model()  # 前向传播
    loss = criterion(outputs, )  # 计算损失
    loss.backward()  # 反向传播
    optimizer.step()  # 更新参数

    if (epoch+1) % 10 == 0:
        print(f'Epoch [{epoch+1}/100], Loss: {loss.item():.4f}')